import joblib
from matplotlib import pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from sklearn import metrics
import numpy as np
from collections import Counter
from utils import draw
from imblearn.over_sampling import RandomOverSampler
from sklearn.model_selection import train_test_split

def random_forest_train(datasets,labels):
    # 训练集和测试集划分
    # X_train = datasets[:115]  # 第0到115个数
    # y_train = labels[:115]
    # X_test = datasets[90:]  # 第90到最后一个数
    # y_test = labels[90:]

    X_train, X_test, y_train, y_test = train_test_split(datasets, labels, test_size=0.2, random_state=40)


    #过采样
    # oversampler = RandomOverSampler()
    # X_train_oversampled, y_train_oversampled = oversampler.fit_resample(X_train, y_train)

    # 随机森林分类器
    clf = RandomForestClassifier(n_estimators=200, random_state=0)
    clf.fit(X_train, y_train)  # 使用训练集对分类器训练
    joblib.dump(clf, "./files/random.pkl")
    y_predict = clf.predict(X_test)  # 使用分类器对测试集进行预测

    auc = metrics.accuracy_score(y_test, y_predict)
    macro = metrics.precision_score(y_test, y_predict, average='macro')
    micro = metrics.precision_score(y_test, y_predict, average='micro')
    macro_recall = metrics.recall_score(y_test, y_predict, average='macro')
    weighted = metrics.f1_score(y_test, y_predict, average='weighted')
    print('准确率:', auc)  # 预测准确率输出
    print('宏平均精确率:', macro)  # 预测宏平均精确率输出
    print('微平均精确率:', micro)  # 预测微平均精确率输出
    print('宏平均召回率:', macro_recall)  # 预测宏平均召回率输出
    print('平均F1-score:', weighted)  # 预测平均f1-score输出
    print('混淆矩阵输出:\n', metrics.confusion_matrix(y_test, y_predict))  # 混淆矩阵输出
    print('分类报告:', metrics.classification_report(y_test, y_predict))  # 分类报告输出
    draw.plot_roc(y_test, y_predict, auc, macro, macro_recall, weighted)  # 绘制ROC曲线并求出AUC值
    return auc, weighted  # 返回准确率与召回率以记录

def random_forest_test(datasets,labels):
    X_test = datasets[:]
    clf0 = joblib.load("./files/random.pkl")
    y_predict = clf0.predict(X_test)  # 使用分类器对测试集进行预测
    np.savetxt('./files/random_result.txt', y_predict)

    # 计算预测正确的百分比并打印
    accuracy = np.mean(y_predict == labels) * 100
    print(f"预测正确的百分比: {accuracy}%")

    #绘制饼状图
    Counter(y_predict)  # {label:sum(label)}
    Yes = sum(y_predict == 1)
    No = sum(y_predict == 0)
    plt.rcParams['font.sans-serif'] = 'SimHei'  # 设置中文显示
    plt.figure(figsize=(6, 6))  # 将画布设定为正方形，则绘制的饼图是正圆
    label = ['有缺陷数', '无缺陷数']  # 定义饼图的标签，标签是列表
    explode = [0.01, 0.05]  # 设定各项距离圆心n个半径
    values = [Yes, No]
    plt.pie(values, explode=explode, labels=label, autopct='%1.1f%%')  # 绘制饼图
    plt.title('缺陷数目')
    plt.show()